import argparse
import simnibs
from simnibs import sim_struct
import nibabel as nib
import os
import platform
import subprocess

def open_itk_snap_cond(filename):
    module = "itk-snap" if platform.system() == 'Windows' else "itksnap"
    try:
        # Use subprocess to run the itk-snap command
        subprocess.Popen([module, '-g', filename])
        print(f"ITK-SNAP opened file: {filename}")
    except subprocess.CalledProcessError as e:
        print(f"Error running ITK-SNAP: {e}")
    except FileNotFoundError:
        print("ITK-SNAP not found. Please ensure it is installed and in your PATH.")
# Function to run the process for each anisotropy type

def process_anisotropy_type(types, subpath, output_path=None):
    # Create session
    S = sim_struct.SESSION()
    S.subpath = subpath
    S.pathfem = os.path.join(subpath, "cond_file")
    # create this file if not exits,
    if not os.path.exists(S.pathfem):
        os.makedirs(S.pathfem)
        print(f"Created directory: {S.pathfem}")
    else:
        print(f"Directory already exists: {S.pathfem}")
    
    tdcs = S.add_tdcslist()
    S._prepare()

    # Set anisotropy type and calculate conductivity
    tdcs.anisotropy_type = types
    cond = tdcs.cond2elmdata(tdcs.mesh)

    # Load reference NIfTI image for affine and voxel information
    reference = os.path.join(subpath, "T1.nii.gz")  # Use subpath to form reference path
    image = nib.load(reference)
    affine = image.affine
    n_voxels = image.header['dim'][1:4]

    # Determine the output file path
    if output_path is None:
        output_path = os.path.join(S.pathfem, f"cond_result_{types}.nii.gz")

    # Convert to NIfTI and save
    cond.to_nifti(n_voxels, affine, fn=output_path, qform=image.header.get_qform(), method='assign')
    print(f"Saved result for {types} anisotropy at {output_path}")

# Main function with argument parsing
def main():
    parser = argparse.ArgumentParser(description="Process FEM anisotropy types and export results to NIfTI")
    parser.add_argument('--subpath', required=True, help="Path to the subject folder (e.g., 'm2m_ernie')")
    parser.add_argument('--output_path', help="Optional: Output file path for the NIfTI file")
    args = parser.parse_args()

    # Process anisotropy types in a loop
    for types in ["scalar", "vn", "dir", "mc"]:
        cond_file_path = os.path.join(args.subpath, "cond_file", f"cond_result_{types}.nii.gz")

        # Check if the file already exists
        if os.path.exists(cond_file_path):
            print(f"File {cond_file_path} already exists. Opening in ITK-SNAP...")
            open_itk_snap_cond(cond_file_path)
        else:
            print(f"File {cond_file_path} does not exist. Generating visualization results for {types} anisotropy...")
            process_anisotropy_type(types, args.subpath, args.output_path)
            open_itk_snap_cond(cond_file_path)

if __name__ == "__main__":
    main()
